{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1: 13s, trn_loss: 1.8947, trn_acc: 37.22%, adv_loss: 1.8952, adv_acc: 37.25%\n",
      "Epoch 2: 16s, trn_loss: 1.0717, trn_acc: 60.74%, adv_loss: 1.0729, adv_acc: 60.69%\n",
      "Epoch 3: 17s, trn_loss: 0.8675, trn_acc: 65.08%, adv_loss: 0.8688, adv_acc: 65.19%\n",
      "Epoch 4: 20s, trn_loss: 0.7730, trn_acc: 71.28%, adv_loss: 0.7744, adv_acc: 71.18%\n",
      "Epoch 5: 13s, trn_loss: 0.6936, trn_acc: 74.45%, adv_loss: 0.6953, adv_acc: 74.42%\n",
      "Epoch 6: 18s, trn_loss: 0.6499, trn_acc: 76.08%, adv_loss: 0.6520, adv_acc: 75.96%\n",
      "Epoch 7: 13s, trn_loss: 0.6133, trn_acc: 78.57%, adv_loss: 0.6157, adv_acc: 78.42%\n",
      "Epoch 8: 15s, trn_loss: 0.5690, trn_acc: 79.84%, adv_loss: 0.5719, adv_acc: 79.56%\n",
      "Epoch 9: 21s, trn_loss: 0.5669, trn_acc: 79.49%, adv_loss: 0.5702, adv_acc: 79.35%\n",
      "Epoch 10: 16s, trn_loss: 0.5311, trn_acc: 81.41%, adv_loss: 0.5347, adv_acc: 81.24%\n",
      "Your estimated attack ability, by applying your attack method on your own trained model, is: 1.2447\n",
      "Your estimated defence ability, by evaluating your own defence model over your attack, is: 0.8034\n",
      "epsilon p:  tensor(0.1000)\n"
     ]
    }
   ],
   "source": [
    "############################################################################\n",
    "### Written by Gaojie Jin and updated by Xiaowei Huang, 2021\n",
    "###\n",
    "### For a 2-nd year undergraduate student competition on \n",
    "### the robustness of deep neural networks, where a student \n",
    "### needs to develop \n",
    "### 1. an attack algorithm, and \n",
    "### 2. an adversarial training algorithm\n",
    "###\n",
    "### The score is based on both algorithms. \n",
    "############################################################################\n",
    "\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "from torchvision import transforms\n",
    "from torch.autograd import Variable\n",
    "import argparse\n",
    "import time\n",
    "import copy\n",
    "\n",
    "# input id\n",
    "id_ = 1000\n",
    "\n",
    "# setup training parameters\n",
    "parser = argparse.ArgumentParser(description='PyTorch MNIST Training')\n",
    "parser.add_argument('--batch-size', type=int, default=128, metavar='N',\n",
    "                    help='input batch size for training (default: 128)')\n",
    "parser.add_argument('--test-batch-size', type=int, default=128, metavar='N',\n",
    "                    help='input batch size for testing (default: 128)')\n",
    "parser.add_argument('--epochs', type=int, default=10, metavar='N',\n",
    "                    help='number of epochs to train')\n",
    "parser.add_argument('--lr', type=float, default=0.01, metavar='LR',\n",
    "                    help='learning rate')\n",
    "parser.add_argument('--no-cuda', action='store_true', default=False,\n",
    "                    help='disables CUDA training')\n",
    "parser.add_argument('--seed', type=int, default=1, metavar='S',\n",
    "                    help='random seed (default: 1)')\n",
    "\n",
    "\n",
    "args = parser.parse_args(args=[]) \n",
    "\n",
    "# judge cuda is available or not\n",
    "use_cuda = not args.no_cuda and torch.cuda.is_available()\n",
    "#device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
    "device = torch.device(\"cpu\")\n",
    "\n",
    "torch.manual_seed(args.seed)\n",
    "kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}\n",
    "\n",
    "############################################################################\n",
    "################    don't change the below code    #####################\n",
    "############################################################################\n",
    "train_set = torchvision.datasets.FashionMNIST(root='data', train=True, download=True, transform=transforms.Compose([transforms.ToTensor()]))\n",
    "train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True)\n",
    "\n",
    "test_set = torchvision.datasets.FashionMNIST(root='data', train=False, download=True, transform=transforms.Compose([transforms.ToTensor()]))\n",
    "test_loader = DataLoader(test_set, batch_size=args.batch_size, shuffle=True)\n",
    "\n",
    "# define fully connected network\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.fc1 = nn.Linear(28*28, 128)\n",
    "        self.fc2 = nn.Linear(128, 64)\n",
    "        self.fc3 = nn.Linear(64, 32)\n",
    "        self.fc4 = nn.Linear(32, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.fc2(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.fc3(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.fc4(x)\n",
    "        output = F.log_softmax(x, dim=1)\n",
    "        return output\n",
    "\n",
    "##############################################################################\n",
    "#############    end of \"don't change the below code\"   ######################\n",
    "##############################################################################\n",
    "\n",
    "#generate adversarial data, you can define your adversarial method\n",
    "def adv_attack(model, X, y, device):\n",
    "    X_adv = Variable(X.data)\n",
    "    \n",
    "    ################################################################################################\n",
    "    ## Note: below is the place you need to edit to implement your own attack algorithm\n",
    "    ################################################################################################\n",
    "    \n",
    "    random_noise = torch.FloatTensor(*X_adv.shape).uniform_(-0.1, 0.1).to(device)\n",
    "    X_adv = Variable(X_adv.data + random_noise)\n",
    "    \n",
    "    ################################################################################################\n",
    "    ## end of attack method\n",
    "    ################################################################################################\n",
    "    \n",
    "    return X_adv\n",
    "\n",
    "#train function, you can use adversarial training\n",
    "def train(args, model, device, train_loader, optimizer, epoch):\n",
    "    model.train()\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        data = data.view(data.size(0),28*28)\n",
    "        \n",
    "        #use adverserial data to train the defense model\n",
    "        #adv_data = adv_attack(model, data, target, device=device)\n",
    "        \n",
    "        #clear gradients\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        #compute loss\n",
    "        #loss = F.nll_loss(model(adv_data), target)\n",
    "        loss = F.nll_loss(model(data), target)\n",
    "        \n",
    "        #get gradients and update\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "\n",
    "\n",
    "#predict function\n",
    "def eval_test(model, device, test_loader):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            data = data.view(data.size(0),28*28)\n",
    "            output = model(data)\n",
    "            test_loss += F.nll_loss(output, target, size_average=False).item()\n",
    "            pred = output.max(1, keepdim=True)[1]\n",
    "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "    test_accuracy = correct / len(test_loader.dataset)\n",
    "    return test_loss, test_accuracy\n",
    "\n",
    "def eval_adv_test(model, device, test_loader):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            data = data.view(data.size(0),28*28)\n",
    "            adv_data = adv_attack(model, data, target, device=device)\n",
    "            output = model(adv_data)\n",
    "            test_loss += F.nll_loss(output, target, size_average=False).item()\n",
    "            pred = output.max(1, keepdim=True)[1]\n",
    "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "    test_accuracy = correct / len(test_loader.dataset)\n",
    "    return test_loss, test_accuracy\n",
    "\n",
    "#main function, train the dataset and print train loss, test loss for each epoch\n",
    "def train_model():\n",
    "    model = Net().to(device)\n",
    "    \n",
    "    ################################################################################################\n",
    "    ## Note: below is the place you need to edit to implement your own training algorithm\n",
    "    ##       You can also edit the functions such as train(...). \n",
    "    ################################################################################################\n",
    "    \n",
    "    optimizer = optim.SGD(model.parameters(), lr=args.lr)\n",
    "    for epoch in range(1, args.epochs + 1):\n",
    "        start_time = time.time()\n",
    "        \n",
    "        #training\n",
    "        train(args, model, device, train_loader, optimizer, epoch)\n",
    "        \n",
    "        #get trnloss and testloss\n",
    "        trnloss, trnacc = eval_test(model, device, train_loader)\n",
    "        advloss, advacc = eval_adv_test(model, device, train_loader)\n",
    "        \n",
    "        #print trnloss and testloss\n",
    "        print('Epoch '+str(epoch)+': '+str(int(time.time()-start_time))+'s', end=', ')\n",
    "        print('trn_loss: {:.4f}, trn_acc: {:.2f}%'.format(trnloss, 100. * trnacc), end=', ')\n",
    "        print('adv_loss: {:.4f}, adv_acc: {:.2f}%'.format(advloss, 100. * advacc))\n",
    "        \n",
    "    adv_tstloss, adv_tstacc = eval_adv_test(model, device, test_loader)\n",
    "    print('Your estimated attack ability, by applying your attack method on your own trained model, is: {:.4f}'.format(1/adv_tstacc))\n",
    "    print('Your estimated defence ability, by evaluating your own defence model over your attack, is: {:.4f}'.format(adv_tstacc))\n",
    "    ################################################################################################\n",
    "    ## end of training method\n",
    "    ################################################################################################\n",
    "    \n",
    "    #save the model\n",
    "    torch.save(model.state_dict(), str(id_)+'.pt')\n",
    "    return model\n",
    "\n",
    "#compute perturbation distance\n",
    "def p_distance(model, train_loader, device):\n",
    "    p = []\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        data = data.view(data.size(0),28*28)\n",
    "        data_ = copy.deepcopy(data.data)\n",
    "        adv_data = adv_attack(model, data, target, device=device)\n",
    "        p.append(torch.norm(data_-adv_data, float('inf')))\n",
    "    print('epsilon p: ',max(p))\n",
    "\n",
    "    \n",
    "################################################################################################\n",
    "## Note: below is for testing/debugging purpose, please comment them out in the submission file\n",
    "################################################################################################\n",
    "    \n",
    "#Comment out the following command when you do not want to re-train the model\n",
    "#In that case, it will load a pre-trained model you saved in train_model()\n",
    "model = train_model()\n",
    "\n",
    "#Call adv_attack() method on a pre-trained model'\n",
    "#the robustness of the model is evaluated against the infinite-norm distance measure\n",
    "#important: MAKE SURE the infinite-norm distance (epsilon p) less than 0.11 !!!\n",
    "p_distance(model, train_loader, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
